
import torch.nn as nn
from model.extractors.Features_SIFT import Features_SIFT,Features_SIFT_WiITH_MASK
from model.extractors.Features_SP import Features_SP
from model.extractors.Features_SFD2 import Features_SFD2
from model.extractors.Features_3DTile import Features_3DTile

def get_fve(args):
    method = args.fve.method
    return get_fve_method(method,args)
    
def get_fve_method(method,args):

    if method == 'SP_3DTILE':
        backbone = Features_3DTile(args)
        return backbone


    if method == "SIFT":
        backbone = Features_SIFT(args)
        return backbone
    elif method == "SuperPoint" or method == "SP":
        backbone = Features_SP(args)
        return backbone
    elif method == "SFD2":
        backbone = Features_SFD2(args)
        return backbone
    elif method == "SIFT_WiITH_MASK":
        backbone = Features_SIFT_WiITH_MASK(args)
        return backbone
    else:
        return None

class FeatureVectorExtractor(nn.Module):
    """The used networks are composed of a backbone and an aggregation layer.
    """
    def __init__(self, args):
        super().__init__()
        self.backbone = get_fve(args)
        # self.arch_name = args.backbone

    def forward(self, x):
        x = self.backbone(x)
        return x

